{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ce73a887",
   "metadata": {},
   "source": [
    "## Results for the toy experiment with synthetic datasets and analytical solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "876b3969-5115-445c-b0ee-d8b02c3a1513",
   "metadata": {},
   "source": [
    "To familiarize yourself with the two syntetic datasets, take a look at the ```VisualizeDatasets``` notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "97d431a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "import os\n",
    "os.chdir(\"..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eabf4eb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d6537e03-9314-40cc-94ed-b44239ea3db0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dsar/miniconda3/envs/tobiasdisentangle/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:526: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/dsar/miniconda3/envs/tobiasdisentangle/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:527: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/dsar/miniconda3/envs/tobiasdisentangle/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:528: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/dsar/miniconda3/envs/tobiasdisentangle/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:529: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/dsar/miniconda3/envs/tobiasdisentangle/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:530: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/dsar/miniconda3/envs/tobiasdisentangle/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:535: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "from data.four_bars import FourBars, ColorBar "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "329d1b61-e188-4bcf-924c-e3f143e3c8cc",
   "metadata": {},
   "source": [
    "A helper function to sample suitable distortions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "23320e7e-b36c-445e-8826-4e3c064a6b97",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_scale_distort(ndim=3, max_val=10):\n",
    "    \"\"\" Sample a distortion matrix A that keeps the factors in the valid range [0, 10] \"\"\"\n",
    "    while True:\n",
    "        ent_matrix = torch.randn(ndim,ndim)\n",
    "        ent_matrix_norm = ent_matrix / ent_matrix.norm(dim=0, keepdim=True)\n",
    "        A = torch.norm(torch.eig(ent_matrix_norm)[0], dim=1)\n",
    "        if torch.min(A) > 0.3 and torch.max(A) < 3.0:\n",
    "            print(A)\n",
    "            break\n",
    "    return ent_matrix_norm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c9162b-dc4b-4d70-8f52-d332408d7785",
   "metadata": {},
   "source": [
    "A helper function to compute the gradients of the faithful encoder distorted by the matrix $M$ via the decoder. \n",
    "According to the transfer lemma stated in the paper, we know that\n",
    "\n",
    "$$J_f(g(z))^\\top J_g(z) = I$$\n",
    "\n",
    "so we can compute $J_f$ from our knowledge of $J_g(z)$. $J_g$ is computed via torch, as the generator is differentiable in pytorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "119ade33-99bb-4132-bd45-588b4fc61749",
   "metadata": {},
   "outputs": [],
   "source": [
    "from common.attributions import generator_jacobian\n",
    "class GTGenWrapper():\n",
    "    def __init__(self, gt_gen):\n",
    "        self.my_gt = gt_gen\n",
    "    def decode(self, latent):\n",
    "        #print(latent.dtype)\n",
    "        return self.my_gt.sample_observations_from_factors(latent, ret_torch=True)\n",
    "    \n",
    "def compute_mde_gradients(facts_org, m, syn_dataset):\n",
    "    \"\"\" Compute the gradients of the faithful encoder. \n",
    "        facts: latent variables values [num_samples, num_facts]\n",
    "        As f = M f* they are M*grad(f*)\n",
    "    \"\"\"\n",
    "    batch = 64\n",
    "    grad_list = []\n",
    "    for i in range(0, len(facts_org), batch):\n",
    "        grad_matrix = generator_jacobian(GTGenWrapper(syn_dataset), facts_org[i:i+64]) # [B, Z, H, W, C]\n",
    "        grad_list.append(grad_matrix)\n",
    "\n",
    "    grad_matrix = torch.cat(grad_list, dim=0)\n",
    "    grad_matrix_size = grad_matrix.shape\n",
    "    #print(grad_matrix_size)\n",
    "    ## Multipy with m.\n",
    "    jg = grad_matrix.reshape(grad_matrix_size[0], grad_matrix_size[1], -1) # [B, Z, N]\n",
    "    jf = jg/jg.pow(2).sum(dim=2, keepdim=True) # Denormalize to have jg*jf=I\n",
    "    # Find min.\n",
    "    min_index = torch.argmin(torch.min(jg.pow(2).sum(dim=2), dim=1)[0])\n",
    "    #print(jg.pow(2).sum(dim=2)[min_index])\n",
    "    #print(facts_org[min_index])\n",
    "    grad_matrix_multiply = jf.transpose(0,1).reshape(grad_matrix_size[1], -1) #[Z, N]\n",
    "    # Invert norming.\n",
    "    #print(grad_matrix_multiply.shape)\n",
    "    grad_matrix_multiply = m.matmul(grad_matrix_multiply).reshape([grad_matrix_size[1], grad_matrix_size[0]] + list(grad_matrix_size[2:])).transpose(0,1)\n",
    "    return grad_matrix_multiply"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9643834-b9bc-47f1-a67c-483c20a5b55f",
   "metadata": {},
   "source": [
    "After stating some helpers, we can now start implementing the analytical solutions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "5a5f1598-ec7a-4526-9e6b-5d6c7ca04bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" We implement the analytical solutions here in a batch-wise fashion. \"\"\"\n",
    "\n",
    "### HELPER FUNCTIONS\n",
    "def build_pair_indices(n):\n",
    "    list1 = []\n",
    "    list2 = []\n",
    "    for i in range(n):\n",
    "        list1.append(i*torch.ones(i,dtype=torch.int64))\n",
    "        list2.append(torch.arange(i, dtype=torch.int64))\n",
    "    return torch.cat(list1), torch.cat(list2)\n",
    "\n",
    "def find_non_sing_submatrix(f):\n",
    "    \"\"\" Return a non singular, square submatrix of M.\"\"\"\n",
    "    #print(f.shape)\n",
    "    # 1st step: Find non singular submatrix.\n",
    "    ind_set = []\n",
    "    clm_count = 0\n",
    "    for i in range(0, len(f)):\n",
    "        ind_set.append(clm_count)\n",
    "        while torch.svd(f[:,ind_set], compute_uv=False)[1][-1] < 1e-5: # linear independence check. Continue if independent.\n",
    "            clm_count += 1\n",
    "            ind_set[-1] = clm_count\n",
    "        #print(\"Appending column\", clm_count, \"to set.\")\n",
    "        clm_count += 1\n",
    "        if len(ind_set) == f.size(0):\n",
    "            break\n",
    "    #print(ind_set)\n",
    "    if len(ind_set) < f.size(0):\n",
    "        return None\n",
    "    else:\n",
    "        return f[:, ind_set]\n",
    "    \n",
    "def find_solution_ima(orth_1, orth_2):\n",
    "    \"\"\" The solution according to IMA theory. \n",
    "        Let orth_1 be Sigma(z_1) and orth_2 be Sigma(z_2) as stated in App. B 7.2. in our paper.\n",
    "    \"\"\"\n",
    "    U = torch.inverse(torch.cholesky(orth_1)) \n",
    "    Gamma = U @ orth_2 @ U.transpose(-2, -1) \n",
    "    # Compute eigentvectors of Gamma, batchwise computation due to old version of torch.\n",
    "    ds_list, q_list = [], []\n",
    "    for i in range(len(orth_1)):\n",
    "        Dslash, Q = torch.symeig(Gamma[i], eigenvectors=True)\n",
    "        ds_list.append(Dslash), q_list.append(Q)\n",
    "    Q = torch.stack(q_list)\n",
    "    return Q.transpose(-2, -1) @ U\n",
    "\n",
    "def find_solution_dma(jf):\n",
    "    \"\"\" Find solution with DMA, using the jacobians. \"\"\"\n",
    "    sol_list = []\n",
    "    for i in range(len(jf)):\n",
    "        sol_list.append(torch.inverse(find_non_sing_submatrix(jf[i])))\n",
    "    Q = torch.stack(sol_list)\n",
    "    return Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "18c13520-61ab-4a27-bea4-419dc0176af6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\" Compute the DCI scores via their entropy formulation: Normalize rows and then compute column-wise entropies.\"\"\"\n",
    "def dci_ent(prod, n_dims = 4):\n",
    "    if len(prod.shape) == 2:\n",
    "        prod = prod.unsqueeze(0)\n",
    "    score_list = []\n",
    "    for i in range(len(prod)):\n",
    "        prodi = torch.abs(prod[i]) / torch.sum(torch.abs(prod[i]), dim=0)\n",
    "        #print(prod)\n",
    "        score = torch.sum(-torch.log(prodi+1e-8)*prodi, dim=0)\n",
    "        score_list.append(1.0-(torch.mean(score)/np.log(n_dims)))\n",
    "    return np.array(score_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "1cae7562-f646-4be0-a904-3fc95b63c0d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Use this flag to switch between datasets.\n",
    "use_colorbar = False\n",
    "n_intervals = 11\n",
    "if use_colorbar:\n",
    "    syn_dataset = ColorBar(n_intervals, nonlin_colors=True)\n",
    "    num_factors = 3\n",
    "else:\n",
    "    syn_dataset = FourBars(n_intervals)\n",
    "    num_factors = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "c5d58875-e7b5-4ba1-bfc9-851fb9028dc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "n_runs = 5\n",
    "n_samples = 20 # How many samples to draw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "96bc1b51-a26d-444a-a870-aee9102b6715",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting run  0\n",
      "tensor([0.8417, 0.8740, 0.8740, 0.4231])\n",
      "Starting run  1\n",
      "tensor([0.9611, 1.2589, 0.5646, 0.5646])\n",
      "Starting run  2\n",
      "tensor([0.7611, 0.7877, 0.7877, 1.1722])\n",
      "Starting run  3\n",
      "tensor([1.1274, 0.4057, 0.4057, 1.0892])\n",
      "Starting run  4\n",
      "tensor([0.5171, 0.5171, 1.2037, 0.6642])\n"
     ]
    }
   ],
   "source": [
    "res_ima_list = []\n",
    "res_dma_list = []\n",
    "for run in range(n_runs):\n",
    "    print(\"Starting run \", run)\n",
    "    facts= torch.rand(n_samples, num_factors)*10\n",
    "    A = sample_scale_distort(num_factors)\n",
    "    grad_matrix = compute_mde_gradients(facts.float(), A, syn_dataset)\n",
    "    grad_matrix = grad_matrix.reshape(len(grad_matrix), grad_matrix.size(1), -1).clone()\n",
    "    orthogonality = torch.bmm(grad_matrix, grad_matrix.transpose(1,2))\n",
    "    \n",
    "    ## IMA (requires gradients at two points, test all combinations of samples against each other)\n",
    "    res1, res2 = build_pair_indices(n_runs)\n",
    "    res =  find_solution_ima(orthogonality[res1], orthogonality[res2])\n",
    "    prod = res@A\n",
    "    res_ima_list.append(dci_ent(prod, num_factors).mean())\n",
    "    \n",
    "    ## DMA (test only one point)\n",
    "    res_dma = find_solution_dma(grad_matrix)\n",
    "    prod = res_dma@A\n",
    "    res_dma_list.append(dci_ent(prod, num_factors).mean())\n",
    "\n",
    "res_dma_list = np.array(res_dma_list)\n",
    "res_ima_list = np.array(res_ima_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "0c119ea8-6d31-44c2-b4ce-f018f14c9177",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IMA-DCI: 0.14445782 +- 0.030627714\n",
      "DMA-DCI: 0.9999982 +- 9.344822e-07\n"
     ]
    }
   ],
   "source": [
    "print(\"IMA-DCI:\", res_ima_list.mean(), \"+-\", res_ima_list.std())\n",
    "print(\"DMA-DCI:\", res_dma_list.mean(), \"+-\", res_dma_list.std())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5b3cf60-fd6d-4452-9e61-b7eac455ee95",
   "metadata": {},
   "source": [
    "We see that, when using ```ColorBar```, we obtain the perfect score with the IMA method, whereas for ```FourBars```, the DMA method yields DCI=1."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec1e42e6-8fc2-4876-960a-d829664c16db",
   "metadata": {},
   "source": [
    "Note: Scores for non-working methods are highly volatile depending on the matrices sampled, and therefore do not exactly match the numbers given in the paper."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "disentangle",
   "language": "python",
   "name": "disentangle"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
